package com.etl.dataflow.scheduler.timewheel;

import com.etl.dataflow.scheduler.cron.CronExpression;

import java.util.*;
import java.util.concurrent.*;

/**
 * 时间轮：月轮
 *
 * @author dx
 * @since 2022/6/17
 */
public class TimeWheel {

    private static final long ERROR_MILL = 10L;
    private final Map<Integer, Queue<WheelNode>> wheel;
    private final WheelType wheelType;
    private final ExecutorService executorService = Executors.newCachedThreadPool();
    private final ScheduledExecutorService scheduledExecutorService = new ScheduledThreadPoolExecutor(1);
    private TimeWheel parent;
    private TimeWheel child;

    public TimeWheel(WheelType wheelType, int stickNum) {
        this.wheelType = wheelType;
        wheel = new HashMap<>((int) Math.ceil(stickNum / 0.75));
        this.start();
    }

    private void start() {
        {
            executorService.execute(() -> scheduledExecutorService.scheduleAtFixedRate(() -> {
                int wheelIndex = this.getCurrentWheelIndex();
                Queue<WheelNode> wheelNodes = this.wheel.get(wheelIndex);
                if (wheelNodes != null && !wheelNodes.isEmpty()) {
                    Queue<WheelNode> consumedNodes = new ArrayDeque<>();
                    WheelNode node;
                    while (!wheelNodes.isEmpty()) {
                        node = wheelNodes.poll();
                        executorService.execute(node.getTask());
                        consumedNodes.add(node);
                    }
                    // 重新计算下次执行时间放入时间轮
                    executorService.submit(() -> {
                        Date date = new Date();
                        Date nextValidTimeAfter;
                        WheelNode newInstance;
                        for (WheelNode consumedNode : consumedNodes) {
                            nextValidTimeAfter = consumedNode.getCronExpression().getNextValidTimeAfter(date);
                            newInstance = consumedNode.createNewInstance(nextValidTimeAfter);
                            this.addTask(newInstance.getCronExpression(), newInstance.getTask());
                        }
                    });
                }

            }, 0L, 1, TimeUnit.SECONDS));
        }
    }

    private int getCurrentWheelIndex() {
        Calendar instance = Calendar.getInstance();
        if (instance.getTimeInMillis() > ERROR_MILL) {
            try {
                Thread.sleep(1);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        return instance.get(Calendar.HOUR_OF_DAY) * 3600 + instance.get(Calendar.MINUTE) * 60 + instance.get(Calendar.SECOND);
    }

    public void addTask(CronExpression cronExpression, WheelTask task) {
        assert task != null;
        TimeExpression timeExpression = TimeExpression.of(cronExpression.getNextValidTimeAfter(new Date()));
        WheelType targetWheelType;
        if (timeExpression.isCurrentDay()) {
            targetWheelType = WheelType.DAY_WHEEL;
        } else if (timeExpression.isCurrentMonth()) {
            targetWheelType = WheelType.MONTH_WHEEL;
        } else {
            targetWheelType = WheelType.YEAR_WHEEL;
        }
        if (targetWheelType.ordinal() == this.wheelType.ordinal()) {
            add(timeExpression, cronExpression, task);
        } else if (targetWheelType.ordinal() < this.wheelType.ordinal()) {
            this.parent.addTask(cronExpression, task);
        } else {
            this.child.addTask(cronExpression, task);
        }
    }

    private void add(TimeExpression timeExpression, CronExpression cronExpression, WheelTask task) {
        int wheelIndex;
        Queue<WheelNode> nodeQueue;
        switch (this.wheelType) {
            case DAY_WHEEL:
                wheelIndex = timeExpression.getHour() * 3600 + timeExpression.getMinute() * 60 + timeExpression.getSecond();
                break;
            case MONTH_WHEEL:
                wheelIndex = timeExpression.getDay();
                break;
            default:
                wheelIndex = timeExpression.getMonth();
                break;
        }
        nodeQueue = wheel.get(wheelIndex);
        if (nodeQueue == null) {
            nodeQueue = new ArrayDeque<>();
            wheel.put(wheelIndex, nodeQueue);
        }
        nodeQueue.add(WheelNode.of(task, timeExpression, cronExpression));
    }

    public void removeTask(String wheelTaskId) {
        Map<WheelType, Boolean> visited = new HashMap<>(8);
        this.removeTask(wheelTaskId, visited);
    }

    private void removeTask(String wheelTaskId, Map<WheelType, Boolean> visited) {
        visited.put(this.wheelType, true);
        for (Map.Entry<Integer, Queue<WheelNode>> entry : this.wheel.entrySet()) {
            entry.getValue().removeIf(wheelNode -> wheelTaskId.equals(wheelNode.getTask().getId()));
        }
        if (this.parent != null && !visited.containsKey(this.parent.wheelType)) {
            this.parent.removeTask(wheelTaskId, visited);
        }
        if (this.child != null && !visited.containsKey(this.child.wheelType)) {
            this.child.removeTask(wheelTaskId, visited);
        }
    }

    public void setParent(TimeWheel parent) {
        this.parent = parent;
    }

    public void setChild(TimeWheel child) {
        this.child = child;
    }

}
